import torch
import torch.nn as nn
import torch.optim as optim

from data_manager.loaders.data_loader import get_loader
from models_manger import RNNModel
from trainer_manger.trainer_manger import TrainerManger

torch.manual_seed(0)
"""
控制随机初始化
    影响 权重初始化（如 nn.Linear、nn.Conv2d 的初始参数）。
    影响 数据打乱（如 DataLoader(shuffle=True)）。
    影响 Dropout、随机增强（如 transforms.RandomCrop）等涉及随机性的操作。"""


def test(layer=""):
    device = torch.device("cuda")
    model = RNNModel(1, 1, 10, 1, device=device)
    print(model)
    criterion = nn.MSELoss(reduction='none')
    optimizer = optim.Adam(model.parameters(), lr=0.002)

    train_loader, test_loader = get_loader("gs10")

    epoch = 200
    # Adam优化器
    trainer_manger = TrainerManger(model, optimizer, criterion, train_loader, test_loader, epoch, device=device)
    trainer_manger.train(True)


if __name__ == '__main__':
    test()
